Write a custom CUDA kernel to optimize the Jensen-Shannon Divergence (JSD) calculation.

The operation computes the divergence between two probability distributions P and Q.
Formula: JSD(P||Q) = 0.5 * KL(P||M) + 0.5 * KL(Q||M), where M = 0.5 * (P + Q).
This involves element-wise computations: p * log(p/m) + q * log(q/m), followed by a summation along the last dimension.

Problem Analysis:
1. Memory Bound: A standard implementation calculates M (write to global), then logs (write/read), then multiplies and sums. This creates huge intermediate memory traffic.
2. Numerical Sensitivity: The calculation involves logarithms and divisions. Precision issues can arise, necessitating double precision (float64).

Optimization Strategy: Fused Block-Reduction Kernel (Double Precision)

1. Block-per-Sample Parallelism: Launch one thread block per sample (row) of the input distributions.

2. Fused Element-wise Compute: Within the block, threads iterate over the dimension D. Each thread loads p and q, computes m = 0.5(p+q), and immediately computes the partial divergence contribution: 0.5 * (p * log(p/m) + q * log(q/m)). This is done entirely in registers.

3. Safe Math: Handle edge cases where p=0 or q=0 (limit is 0) to avoid NaNs.

4. Shared Memory Reduction: Threads aggregate their partial sums using a tree-based parallel reduction in shared memory to produce the final scalar JSD for that sample.

5. Vectorized Access: Use `double2` (128-bit) loads to maximize global memory throughput. 
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 1024
DIM = 4096
SHAPE = (BATCH_SIZE, DIM)

DTYPE = torch.float64

class JSDivergence(nn.Module):
    """
    Jensen-Shannon Divergence.
    """
    def __init__(self):
        super(JSDivergence, self).__init__()

    def forward(self, p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
        # 1. Calculate Mean Distribution M
        m = 0.5 * (p + q)
        
        # 2. Compute KL(P||M) and KL(Q||M)
        term1 = 0.5 * (p * (p.log() - m.log()))
        term2 = 0.5 * (q * (q.log() - m.log()))
        
        term1[p == 0] = 0.0
        term2[q == 0] = 0.0
        
        return (term1 + term2).sum(dim=-1)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.jsd = JSDivergence()
    
    def forward(self, p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
        return self.jsd(p, q)

def get_inputs():
    p_logits = torch.randn(SHAPE, dtype=DTYPE)
    q_logits = torch.randn(SHAPE, dtype=DTYPE)
    
    p = F.softmax(p_logits, dim=-1)
    q = F.softmax(q_logits, dim=-1)
    
    return [p.contiguous(), q.contiguous()]

def get_init_inputs():
    return []